# service_consumer.py
import random
from functools import wraps
from typing import Optional

import aiohttp
import consul
from fastapi import FastAPI, HTTPException

app = FastAPI()


async def discover_service_with_lb(service_name):
    c = consul.Consul(host='127.0.0.1', port=8500)
    index, services = c.health.service(service_name, passing=True)

    if not services:
        raise HTTPException(status_code=503, detail=f"Service {service_name} not available")

    # 随机选择一个服务实例
    service = random.choice(services)
    return f"http://{service['Service']['Address']}:{service['Service']['Port']}"


def with_service_discovery(service_name):
    def decorator(func):
        @wraps(func)
        async def wrapper(*args, **kwargs):
            service_url = await discover_service_with_lb(service_name)
            kwargs['service_url'] = service_url
            return await func(*args, **kwargs)

        return wrapper

    return decorator


@app.get("/call-provider")
@with_service_discovery("admin-provider")
async def call_provider(service_url: Optional[str] = None):
    async with aiohttp.ClientSession() as session:
        async with session.get(f"{service_url}/user/find_all") as response:
            return await response.json()


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=8001)
